from app.models.backtest import BacktestRun, DailyPnl
from app.strategy.grid_gnn import GridGNNStrategy
from clickhouse_driver import Client
import pandas as pd
import numpy as np
from datetime import date, datetime
from loguru import logger


CH = Client(host='clickhouse', port=9000, database='qbot')

def create_bt_tables():
    CH.execute('''
        CREATE TABLE IF NOT EXISTS backtestrun (
            id         UInt32,
            name       String,
            start_date Date,
            end_date   Date,
            init_cash  Float64,
            status     String,
            created_at DateTime
        ) ENGINE = MergeTree()
        PARTITION BY toYYYYMM(created_at)
        ORDER BY (id, created_at)
    ''')
    CH.execute('''
        CREATE TABLE IF NOT EXISTS daily_pnl (
            run_id    UInt32,
            trade_date Date,
            nav       Float64,
            pnl       Float64,
            drawdown  Float64,
            sharpe    Float64
        ) ENGINE = MergeTree()
        PARTITION BY toYYYYMM(trade_date)
        ORDER BY (run_id, trade_date)
    ''')

def perf_stats(pnl: pd.Series) -> dict:
    """自写绩效：年化、夏普、最大回撤"""
    ret = pnl.pct_change().dropna()
    ann_ret = ret.mean() * 252
    sharpe = ret.mean() / ret.std() * np.sqrt(252) if ret.std() else 0
    cum = (1 + ret).cumprod()
    drawdown = (cum / cum.cummax() - 1).min()
    return {
        "cagr": ann_ret,
        "sharpe": sharpe,
        "max_drawdown": drawdown,
    }

def run_backtest(name: str, codes: list[str], start: date, end: date, init_cash: float = 1e6):
    logger.warning(f"ENTRY | codes={codes} start={start} end={end}")
    run_id = int(datetime.now().timestamp())   # 自增 UInt32
    run = BacktestRun(id=run_id, name=name, start_date=start, end_date=end, init_cash=init_cash, status='running', created_at=datetime.utcnow())
    # 替换这一行
    CH.execute(
        'INSERT INTO backtestrun (id, name, start_date, end_date, init_cash, status, created_at) VALUES',
        [(run.id, run.name, run.start_date, run.end_date, run.init_cash, run.status, run.created_at)]
    )

    strategy = GridGNNStrategy(codes, start, end)
    print("[DEBUG] strategy:", strategy)
    logger.debug(f"strategy: {strategy}")
    logger.debug(f"strategy: {strategy.bars.head()}")
    logger.debug(f"strategy: {strategy.bars.columns.tolist()}")
    bars = strategy.bars.set_index(['date', 'code']).sort_index()
    dates = bars.index.get_level_values(0).unique()

    nav = init_cash
    records = []
    for d in dates:
        day_bars = bars.loc[d]
        weights = strategy.next(day_bars.iloc[0])
        # 简化：按权重买入，收盘净值 = 加权收盘
        port_ret = (day_bars['adj_close'] * pd.Series(weights)).sum() / day_bars['close'].sum() - 1
        nav *= (1 + port_ret)
        records.append({
            'run_id': run.id,
            'date': d,
            'nav': nav,
            'pnl': nav - init_cash,
            'drawdown': (nav / np.maximum.accumulate(records[-1]['nav'] if records else nav) - 1),
        })

    df_rec = pd.DataFrame(records)
    stats = perf_stats(df_rec['pnl'])   # 夏普、年化、最大回撤
    CH.execute('INSERT INTO daily_pnl (run_id, trade_date, nav, pnl, drawdown, sharpe) VALUES', df_rec.itertuples(index=False, name=None))
    CH.execute("UPDATE backtestrun SET status = 'D' WHERE id = %(id)s", {'id': run.id})
    return run.id, stats
